import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, random_split, Subset
import os
import numpy as np


class MNISTLoader:
    """
    MNIST数据集加载器类
    """
    
    def __init__(self, data_dir='./data', use_cuda=True):
        """
        初始化数据加载器
        
        参数:
        - data_dir: 数据保存目录
        - use_cuda: 是否使用CUDA (影响pin_memory等设置)
        """
        self.data_dir = data_dir
        self.use_cuda = use_cuda and torch.cuda.is_available()
        
        # 确保数据目录存在
        os.makedirs(data_dir, exist_ok=True)
        
        # 数据集统计信息
        self.mean = (0.1307,)
        self.std = (0.3081,)
        
        # 初始化数据集和加载器
        self.train_dataset = None
        self.val_dataset = None
        self.test_dataset = None
        self.train_loader = None
        self.val_loader = None
        self.test_loader = None
    
    def get_transforms(self, augment=True, advanced_augmentation=False):
        """
        获取数据转换
        
        参数:
        - augment: 是否使用数据增强
        - advanced_augmentation: 是否使用高级数据增强（仅在augment为True时有效）
        """
        transform_list = []
        
        # MNIST专用预处理
        transform_list.append(transforms.Resize((28, 28)))
        
        if augment:
            # 基本数据增强
            transform_list.extend([
                transforms.RandomRotation(30),
                transforms.RandomAffine(degrees=0,
                translate=(0.3, 0.3),
                scale=(0.5, 2),
                shear=10
                ),
                transforms.ElasticTransform(alpha=20.0, sigma=3.0),
                transforms.RandomPerspective(distortion_scale=0.3, p=0.5),
                transforms.RandomCrop(28, padding=1)
            ])
            
            if advanced_augmentation:
                transform_list.extend([
                    transforms.ElasticTransform(alpha=30.0, sigma=5.0),
                    transforms.RandomApply([transforms.ElasticTransform(alpha=10.0, sigma=1.0)], p=0.1)
                ])
        
        
        transform_list.extend([
            transforms.ToTensor(),
            transforms.Normalize(self.mean, self.std)
        ])
        
        return transforms.Compose(transform_list)
        
    def load_datasets(self, validation_split=0.1, advanced_augmentation=False, use_augmentation=True):
        """
        加载数据集
        
        参数:
        - validation_split: 验证集比例
        - advanced_augmentation: 是否使用高级数据增强
        - use_augmentation: 是否使用数据增强
        """
        # 获取转换，根据use_augmentation决定是否启用数据增强
        train_transform = self.get_transforms(augment=use_augmentation, advanced_augmentation=use_augmentation and advanced_augmentation)
        test_transform = self.get_transforms(augment=False)
        
        # 加载MNIST数据集
        self.train_dataset = datasets.MNIST(
            root=self.data_dir, 
            train=True, 
            download=True, 
            transform=train_transform
        )
        
        self.test_dataset = datasets.MNIST(
            root=self.data_dir, 
            train=False, 
            download=True, 
            transform=test_transform
        )
        
        # 分割训练集为训练集和验证集
        train_size = int((1 - validation_split) * len(self.train_dataset))
        val_size = len(self.train_dataset) - train_size
        
        generator = torch.Generator().manual_seed(42)
        self.train_dataset, self.val_dataset = random_split(
            self.train_dataset, 
            [train_size, val_size],
            generator=generator
        )
        
        # 验证集使用测试集的转换
        self.val_dataset.dataset.transform = test_transform
        
        return self.train_dataset, self.val_dataset, self.test_dataset
    
    def create_dataloaders(self, batch_size=64, num_workers=None, shuffle_train=True, 
                          persistent_workers=True, prefetch_factor=3, pin_memory=None):
        """
        创建数据加载器
        
        参数:
        - batch_size: 批次大小
        - num_workers: 数据加载的并行工作进程数
        - shuffle_train: 是否打乱训练集
        - persistent_workers: 是否保持工作进程存活
        - prefetch_factor: 预取的批次数量
        - pin_memory: 是否将数据固定在内存中
        
        返回:
        - train_loader, val_loader, test_loader
        """
        # 如果未指定，自动确定参数
        if num_workers is None:
            num_workers = min(os.cpu_count(), 8)
        
        if pin_memory is None:
            pin_memory = self.use_cuda
        
        # 创建数据加载器
        loader_kwargs = {
            'num_workers': num_workers,
            'pin_memory': pin_memory,
            'persistent_workers': persistent_workers if num_workers > 0 else False,
        }
        
        # 只有当num_workers > 0时才设置prefetch_factor
        if num_workers > 0:
            loader_kwargs['prefetch_factor'] = prefetch_factor
        
        self.train_loader = DataLoader(
            self.train_dataset,
            batch_size=batch_size,
            shuffle=shuffle_train,
            **loader_kwargs
        )
        
        self.val_loader = DataLoader(
            self.val_dataset,
            batch_size=batch_size * 2,
            shuffle=False,
            **loader_kwargs
        )
        
        self.test_loader = DataLoader(
            self.test_dataset,
            batch_size=batch_size * 2,
            shuffle=False,
            **loader_kwargs
        )
        
        return self.train_loader, self.val_loader, self.test_loader
    
    def get_sample_batch(self, dataset='train', num_samples=16):
        """
        获取样本批次用于可视化
        
        参数:
        - dataset: 数据集类型 ('train', 'val', 'test')
        - num_samples: 样本数量
        
        返回:
        - 样本图像和标签
        """
        dataset_map = {
            'train': self.train_dataset,
            'val': self.val_dataset,
            'test': self.test_dataset
        }
        
        dataset = dataset_map.get(dataset, self.train_dataset)
        
        # 随机选择样本
        indices = torch.randperm(len(dataset))[:num_samples]
        sample_dataset = Subset(dataset, indices)
        
        # 创建临时数据加载器
        sample_loader = DataLoader(sample_dataset, batch_size=num_samples)
        
        # 获取一个批次
        for images, labels in sample_loader:
            return images, labels
    
    def get_class_distribution(self):
        """
        获取训练集和测试集的类别分布
        
        返回:
        - 包含类别分布的字典
        """
        # 获取原始数据集的标签
        train_labels = self.train_dataset.dataset.targets[self.train_dataset.indices]
        val_labels = self.val_dataset.dataset.targets[self.val_dataset.indices]
        test_labels = self.test_dataset.targets
        
        # 计算分布
        train_dist = torch.bincount(train_labels).numpy()
        val_dist = torch.bincount(val_labels).numpy()
        test_dist = torch.bincount(test_labels).numpy()
        
        return {
            'train': train_dist,
            'val': val_dist,
            'test': test_dist
        }


def get_mnist_data(batch_size=64, data_dir='./data/MNIST', use_augmentation=True):
    """
    加载MNIST数据集，支持数据增强
    
    参数:
    - batch_size: 批次大小
    - data_dir: 数据存储目录
    - use_augmentation: 是否使用数据增强
    """
    # 使用数据加载器类，但保持原有API兼容
    loader = MNISTLoader(data_dir=data_dir)
    
    # 加载数据集，正确传递use_augmentation参数
    loader.load_datasets(
        validation_split=0,
        use_augmentation=use_augmentation,
        advanced_augmentation=use_augmentation
    )
    
    # 创建数据加载器（不创建验证集，保持原有行为）
    train_loader, _, test_loader = loader.create_dataloaders(
        batch_size=batch_size,
        num_workers=4,
        shuffle_train=True,
        pin_memory=True,
        persistent_workers=True
    )
    
    print(f"数据集加载完成: 训练集大小={len(loader.train_dataset)}, 测试集大小={len(loader.test_dataset)}")
    print(f"数据增强: {'已启用' if use_augmentation else '未启用'}")
    
    return train_loader, test_loader


def get_mnist_dataloaders(batch_size=64, validation_split=0.1, num_workers=4, shuffle=True, pin_memory=True, use_augmentation=True):
    """
    获取MNIST数据集的数据加载器，支持训练集、验证集和测试集
    
    参数:
    - batch_size: 批次大小
    - validation_split: 验证集比例
    - num_workers: 数据加载的并行工作进程数
    - shuffle: 是否随机打乱数据
    - pin_memory: 是否将数据固定在内存中以加速GPU传输
    - use_augmentation: 是否使用数据增强
    
    返回:
    - train_loader: 训练集数据加载器
    - val_loader: 验证集数据加载器
    - test_loader: 测试集数据加载器
    """
    loader = MNISTLoader()
    loader.load_datasets(
        validation_split=validation_split,
        use_augmentation=use_augmentation
    )
    return loader.create_dataloaders(
        batch_size=batch_size,
        num_workers=num_workers,
        shuffle_train=shuffle,
        pin_memory=pin_memory,
        persistent_workers=True,
        prefetch_factor=3
    )

